package com.ruoyi.common.utils;

import org.apache.commons.codec.binary.Base64;

import javax.crypto.Cipher;
import java.io.ByteArrayOutputStream;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.HashMap;
import java.util.Map;

/**
 * RSA加解密工具类
 */
public class RSAUtil {
    /**
     * 加密算法
     */
    private static final String ALGORITHM = "RSA";
    
    /**
     * 填充方式
     */
    private static final String PADDING_FORMAT = "RSA/ECB/PKCS1Padding";
    
    /**
     * 明文块最大值
     */
    private static final int MAX_ENCRYPT_BLOCK = 117;
    
    /**
     * 密文块最大值
     */
    private static final int MAX_DECRYPT_BLOCK = 128;
    
    /**
     * 字符编码
     */
    public static final String CHARSET_UTF_8 = "UTF-8";
    
    /**
     * 生成密钥对
     *
     * @return 包含公钥和私钥的Map
     * @throws Exception 如果生成密钥对失败
     */
    private static Map<String, String> generateKeyPair() throws Exception {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(ALGORITHM);
        keyPairGenerator.initialize(1024);
        
        KeyPair keyPair = keyPairGenerator.generateKeyPair();
        RSAPublicKey publicKey = (RSAPublicKey) keyPair.getPublic();
        RSAPrivateKey privateKey = (RSAPrivateKey) keyPair.getPrivate();
        
        Map<String, String> keyMap = new HashMap<>(2);
        keyMap.put("RSAPublicKey", Base64.encodeBase64String(publicKey.getEncoded()));
        keyMap.put("RSAPrivateKey", Base64.encodeBase64String(privateKey.getEncoded()));
        
        return keyMap;
    }
    
    /**
     * 使用公钥加密
     *
     * @param text      明文
     * @param publicKey 公钥
     * @param charset   字符编码
     * @return 加密后的字符串（base64）
     * @throws Exception 如果加密失败
     */
    public static String encrypt(String text, String publicKey, String charset) throws Exception {
        byte[] data = text.getBytes(charset);
        byte[] keyBytes = Base64.decodeBase64(publicKey);
        X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(keyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
        Key key = keyFactory.generatePublic(x509KeySpec);
        
        return getCipherText(data, key, Cipher.ENCRYPT_MODE, MAX_ENCRYPT_BLOCK, charset);
    }
    
    /**
     * 使用私钥解密
     *
     * @param ciphertext 密文（base64）
     * @param privateKey 私钥
     * @param charset    字符编码
     * @return 解密后的字符串
     * @throws Exception 如果解密失败
     */
    public static String decrypt(String ciphertext, String privateKey, String charset) throws Exception {
        byte[] data = Base64.decodeBase64(ciphertext);
        byte[] keyBytes = Base64.decodeBase64(privateKey);
        PKCS8EncodedKeySpec pkcs8KeySpec = new PKCS8EncodedKeySpec(keyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
        Key key = keyFactory.generatePrivate(pkcs8KeySpec);
        
        return getCipherText(data, key, Cipher.DECRYPT_MODE, MAX_DECRYPT_BLOCK, charset);
    }
    
    /**
     * 使用私钥加密
     *
     * @param text       明文
     * @param privateKey 私钥
     * @param charset    字符编码
     * @return 加密后的字符串（base64）
     * @throws Exception 如果加密失败
     */
    public static String encryptWithPrivateKey(String text, String privateKey, String charset) throws Exception {
        byte[] data = text.getBytes(charset);
        byte[] keyBytes = Base64.decodeBase64(privateKey);
        PKCS8EncodedKeySpec pkcs8KeySpec = new PKCS8EncodedKeySpec(keyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
        Key key = keyFactory.generatePrivate(pkcs8KeySpec);
        
        return getCipherText(data, key, Cipher.ENCRYPT_MODE, MAX_ENCRYPT_BLOCK, charset);
    }
    
    /**
     * 使用公钥解密
     *
     * @param ciphertext 密文（base64）
     * @param publicKey  公钥
     * @param charset    字符编码
     * @return 解密后的字符串
     * @throws Exception 如果解密失败
     */
    public static String decryptWithPublicKey(String ciphertext, String publicKey, String charset) throws Exception {
        byte[] data = Base64.decodeBase64(ciphertext);
        byte[] keyBytes = Base64.decodeBase64(publicKey);
        X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(keyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
        Key key = keyFactory.generatePublic(x509KeySpec);
        
        return getCipherText(data, key, Cipher.DECRYPT_MODE, MAX_DECRYPT_BLOCK, charset);
    }
    
    /**
     * 获取加密或解密后的文本
     *
     * @param data     待处理的数据
     * @param key      密钥
     * @param mode     模式（加密或解密）
     * @param maxBlock 每次处理的最大数据块大小
     * @param charset  字符编码
     * @return 处理后的文本
     * @throws Exception 如果处理失败
     */
    private static String getCipherText(byte[] data, Key key, int mode, int maxBlock, String charset) throws Exception {
        Cipher cipher = Cipher.getInstance(PADDING_FORMAT);
        cipher.init(mode, key);
        int length = data.length;
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        int offSet = 0;
        byte[] cache;
        int i = 0;
        
        // 对数据分段处理
        while (length - offSet > 0) {
            int len = Math.min(length - offSet, maxBlock);
            cache = cipher.doFinal(data, offSet, len);
            out.write(cache, 0, cache.length);
            
            i++;
            offSet = i * maxBlock;
        }
        
        byte[] cipherBytes = out.toByteArray();
        out.close();
        
        if (mode == Cipher.ENCRYPT_MODE) {
            return Base64.encodeBase64String(cipherBytes); // 返回加密后的base64字符串
        } else {
            return new String(cipherBytes, charset); // 返回解密后的字符串
        }
    }
    
    public static void main(String[] args) {
        try {
            // 创建密钥对
            Map<String, String> keys = generateKeyPair();
            System.out.println("==================================================生成Public Key======================================================");
            String pubKey = keys.get("RSAPublicKey");
            System.out.println(pubKey);
            
            System.out.println();
            
            System.out.println("==================================================生成Private Key======================================================");
            String priKey = keys.get("RSAPrivateKey");
            System.out.println(priKey);
            
            System.out.println();
            
            // 测试加密文本
            String text = "123";
            System.out.println("==================================================公钥加密，私钥解密======================================================");
            String ciphertext = encrypt(text, pubKey, CHARSET_UTF_8);
            System.out.println("明文：" + text);
            System.out.println("加密后文本：" + ciphertext);
            System.out.println("解密后文本：" + decrypt(ciphertext, priKey, CHARSET_UTF_8));
            
            System.out.println();
            
            System.out.println("==================================================私钥加密，公钥解密======================================================");
            ciphertext = encryptWithPrivateKey(text, priKey, CHARSET_UTF_8);
            System.out.println("明文：" + text);
            System.out.println("加密后文本：" + ciphertext);
            System.out.println("解密后文本：" + decryptWithPublicKey(ciphertext, pubKey, CHARSET_UTF_8));
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
